import numpy as np

def get_ms(msg):
    if msg.header.stamp.secs==0 and msg.header.stamp.nsecs==0:
        return int(msg.time_us/1000)
    return msg.header.stamp.secs * 1000 + int(msg.header.stamp.nsecs / (10 ** 6))

def topic2path(topic,index=1):
    datas = str.split(topic, '/')
    assert datas[0] == '' and len(datas) >= 2
    return datas[index]


def bbox3d2corners(bboxes):
    '''
    bboxes: shape=(n, 7)
    return: shape=(n, 8, 3)
           ^ z   x            6 ------ 5
           |   /             / |     / |
           |  /             2 -|---- 1 |
    y      | /              |  |     | |
    <------|o               | 7 -----| 4
                            |/   o   |/
                            3 ------ 0
    x: front, y: left, z: top
    '''
    centers, dims, angles = bboxes[:, :3], bboxes[:, 3:6], bboxes[:, 6]

    # kitti之前的centerz 难道是最大高度？？我们的肯定是中心高度
    # 1.generate bbox corner coordinates, clockwise from minimal point
    # bboxes_corners = np.array([[-0.5, -0.5, 0], [-0.5, -0.5, 1.0], [-0.5, 0.5, 1.0], [-0.5, 0.5, 0.0],
    #                            [0.5, -0.5, 0], [0.5, -0.5, 1.0], [0.5, 0.5, 1.0], [0.5, 0.5, 0.0]],
    #                            dtype=np.float32)
    bboxes_corners = np.array([[-0.5, -0.5, -0.5], [-0.5, -0.5, 0.5], [-0.5, 0.5, 0.5], [-0.5, 0.5, -0.5],
                               [0.5, -0.5, -0.5], [0.5, -0.5, 0.5], [0.5, 0.5, 0.5], [0.5, 0.5, -0.5]],
                              dtype=np.float32)
    bboxes_corners = bboxes_corners[None, :, :] * dims[:, None, :] # (1, 8, 3) * (n, 1, 3) -> (n, 8, 3)

    # 2. rotate around z axis
    rot_sin, rot_cos = np.sin(angles).reshape(len(bboxes),1), np.cos(angles).reshape(len(bboxes),1)
    rot_mat = np.concatenate(
        [rot_cos, rot_sin, np.zeros_like(rot_cos),
         -rot_sin, rot_cos, np.zeros_like(rot_cos),
         np.zeros_like(rot_cos),np.zeros_like(rot_cos), np.ones_like(rot_cos)], axis=1).astype(np.float32).reshape(-1, 3, 3)

    # rot_mat = np.array([[rot_cos, -rot_sin, np.zeros_like(rot_cos)],
    #                     [rot_sin, rot_cos, np.zeros_like(rot_cos)],
    #                     [np.zeros_like(rot_cos), np.zeros_like(rot_cos), np.ones_like(rot_cos)]],
    #                     dtype=np.float32) # (3, 3, n)
    # rot_mat = np.transpose(rot_mat, (2, 1, 0)) # (n, 3, 3)
    bboxes_corners = bboxes_corners @ rot_mat # (n, 8, 3)

    # 3. translate to centers
    bboxes_corners += centers[:, None, :]
    return bboxes_corners

def group_rectangle_vertexs(bboxes_corners):
    '''
    bboxes_corners: shape=(n, 8, 3)
      6 ------- 5                                        5      6
     / |      / |              _                       / |    / |
    2 -|----- 1 |              |    2 - 1    6 - 5    1  4   2  7     6 - 5     7 - 4
    |  |      | |       -----> |    |   | +  |   | +  | /  + | /  +  /   /  +  /   /
    | 7 ------| 4              |    3 - 0    7 - 4    0      3      2 - 1     3 - 0
    |/   o    |/               _
    3 ------- 0
    return: shape=(n, 6, 4, 3)
    '''
    rec1 = np.stack([bboxes_corners[:, 0], bboxes_corners[:, 1], bboxes_corners[:, 3], bboxes_corners[:, 2]],
                    axis=1)  # (n, 4, 3)
    rec2 = np.stack([bboxes_corners[:, 4], bboxes_corners[:, 7], bboxes_corners[:, 6], bboxes_corners[:, 5]],
                    axis=1)  # (n, 4, 3)
    rec3 = np.stack([bboxes_corners[:, 0], bboxes_corners[:, 4], bboxes_corners[:, 5], bboxes_corners[:, 1]],
                    axis=1)  # (n, 4, 3)
    rec4 = np.stack([bboxes_corners[:, 2], bboxes_corners[:, 6], bboxes_corners[:, 7], bboxes_corners[:, 3]],
                    axis=1)  # (n, 4, 3)
    rec5 = np.stack([bboxes_corners[:, 1], bboxes_corners[:, 5], bboxes_corners[:, 6], bboxes_corners[:, 2]],
                    axis=1)  # (n, 4, 3)
    rec6 = np.stack([bboxes_corners[:, 0], bboxes_corners[:, 3], bboxes_corners[:, 7], bboxes_corners[:, 4]],
                    axis=1)  # (n, 4, 3)
    group_rectangle_vertexs = np.stack([rec1, rec2, rec3, rec4, rec5, rec6], axis=1)
    return group_rectangle_vertexs

def group_rectangle_vertexs(bboxes_corners):
    '''
    bboxes_corners: shape=(n, 8, 3)
    return: shape=(n, 6, 4, 3)
      6 ------- 5                                        5      6
     / |      / |              _                       / |    / |
    2 -|----- 1 |              |    2 - 1    6 - 5    1  4   2  7     6 - 5     7 - 4
    |  |      | |       -----> |    |   | +  |   | +  | /  + | /  +  /   /  +  /   /
    | 7 ------| 4              |    3 - 0    7 - 4    0      3      2 - 1     3 - 0
    |/   o    |/               _
    3 ------- 0
    '''
    rec1 = np.stack([bboxes_corners[:, 0], bboxes_corners[:, 1], bboxes_corners[:, 3], bboxes_corners[:, 2]], axis=1) # (n, 4, 3)
    rec2 = np.stack([bboxes_corners[:, 4], bboxes_corners[:, 7], bboxes_corners[:, 6], bboxes_corners[:, 5]], axis=1) # (n, 4, 3)
    rec3 = np.stack([bboxes_corners[:, 0], bboxes_corners[:, 4], bboxes_corners[:, 5], bboxes_corners[:, 1]], axis=1) # (n, 4, 3)
    rec4 = np.stack([bboxes_corners[:, 2], bboxes_corners[:, 6], bboxes_corners[:, 7], bboxes_corners[:, 3]], axis=1) # (n, 4, 3)
    rec5 = np.stack([bboxes_corners[:, 1], bboxes_corners[:, 5], bboxes_corners[:, 6], bboxes_corners[:, 2]], axis=1) # (n, 4, 3)
    rec6 = np.stack([bboxes_corners[:, 0], bboxes_corners[:, 3], bboxes_corners[:, 7], bboxes_corners[:, 4]], axis=1) # (n, 4, 3)
    group_rectangle_vertexs = np.stack([rec1, rec2, rec3, rec4, rec5, rec6], axis=1)
    return group_rectangle_vertexs

def group_plane_equation(bbox_group_rectangle_vertexs):
    '''
    bbox_group_rectangle_vertexs: shape=(n, 6, 4, 3)
    return: shape=(n, 6, 4)

      6 ------- 5                                        5      6
     / |      / |              _                       / |    / |
    2 -|----- 1 |              |    2 - 1    6 - 5    1  4   2  7     6 - 5     7 - 4
    |  |      | |       -----> |    |   | +  |   | +  | /  + | /  +  /   /  +  /   /
    | 7 ------| 4              |    3 - 0    7 - 4    0      3      2 - 1     3 - 0
    |/   o    |/               _
    3 ------- 0

    '''
    # 1. generate vectors for a x b
    vectors = bbox_group_rectangle_vertexs[:, :, :2] - bbox_group_rectangle_vertexs[:, :, 1:3]
    normal_vectors = np.cross(vectors[:, :, 0], vectors[:, :, 1]) # (n, 6, 3) 利用边向量的叉乘求平面的法向量 提前設置順序使得法相量指向box外
    normal_d = np.einsum('ijk,ijk->ij', bbox_group_rectangle_vertexs[:, :, 0], normal_vectors) # (n, 6)
    plane_equation_params = np.concatenate([normal_vectors, -normal_d[:, :, None]], axis=-1)
    return plane_equation_params

def points_in_bboxes(points, plane_equation_params):
    N, n = len(points), len(plane_equation_params)
    m = plane_equation_params.shape[1]

    # 初始化一个N x n的全True掩码
    masks = np.ones((N, n), dtype=np.bool_)

    # 抽取plane_equation_params的各项
    A = plane_equation_params[:, :, 0]  # (n, m) -> a系数
    B = plane_equation_params[:, :, 1]  # (n, m) -> b系数
    C = plane_equation_params[:, :, 2]  # (n, m) -> c系数
    D = plane_equation_params[:, :, 3]  # (n, m) -> d系数

    # 提取points的x, y, z
    x = points[:, 0]  # (N,)
    y = points[:, 1]  # (N,)
    z = points[:, 2]  # (N,)

    # 使用np.newaxis扩展x, y, z到 (N, 1) -> (N, n, m)
    x = x[:, np.newaxis, np.newaxis]  # (N, 1, 1)
    y = y[:, np.newaxis, np.newaxis]  # (N, 1, 1)
    z = z[:, np.newaxis, np.newaxis]  # (N, 1, 1)

    # 将x, y, z广播到与平面参数A, B, C, D匹配的形状
    x = np.broadcast_to(x, (N, n, m))  # (N, n, m)
    y = np.broadcast_to(y, (N, n, m))  # (N, n, m)
    z = np.broadcast_to(z, (N, n, m))  # (N, n, m)

    # 计算所有平面方程的值 a*x + b*y + c*z + d
    plane_values = A * x + B * y + C * z + D  # (N, n, m)

    # 找到任何一个平面方程值 >= 0，设置为False
    masks = np.all(plane_values < 0, axis=2)  # (N, n)，沿着m维度检查是否所有值都<0

    return masks